import json
import jieba  # 中文分词库
from simhash import Simhash  # 核心算法库


# ----------------------
# 第一步：数据加载与预处理
# ----------------------
def load_jsonl(file_path):
    """读取JSONL文件，逐行加载数据"""
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            data.append(json.loads(line))
    return data


# ----------------------
# 第二步：特征提取函数
# ----------------------
def get_features(text):
    """生成文本特征：分词 -> 过滤停用词 -> 权重计算"""
    words = jieba.cut(text)
    features = {}
    for word in words:
        if word.strip():  # 去除空格等无效字符
            features[word] = features.get(word, 0) + 1
    return features


# ----------------------
# 第三步：生成SimHash指纹
# ----------------------
def generate_simhash(messages):
    """从对话数据生成唯一标识码"""
    # 拼接所有对话内容
    full_text = ' '.join([
        msg['content'].strip()
        for msg in messages
        if msg['content'].strip()
    ])

    # 生成特征并计算哈希
    features = get_features(full_text)
    return Simhash(features).value


# ----------------------
# 第四步：主处理逻辑
# ----------------------
def deduplicate(data, threshold=2):
    """
    去重核心算法
    :param data: 原始数据集
    :param threshold: 海明距离阈值（默认3，越小越严格）
    :return: 去重后的数据
    """
    hashes = {}  # 存储哈希值与数据索引
    duplicates = set()  # 记录需要删除的索引

    for idx, item in enumerate(data):
        current_hash = generate_simhash(item['messages'])

        # 与已有哈希值比对
        for exist_hash, exist_idx in hashes.items():
            # 计算两个哈希值的二进制差异位数
            hamming_dist = bin(current_hash ^ exist_hash).count('1')

            if hamming_dist <= threshold:
                # 保留较长对话（防止数据截断）
                if len(str(item)) > len(str(data[exist_idx])):
                    duplicates.add(exist_idx)
                else:
                    duplicates.add(idx)
                break

        # 记录当前哈希值
        if current_hash not in hashes:
            hashes[current_hash] = idx

    # 过滤重复项
    return [item for idx, item in enumerate(data) if idx not in duplicates]


# ----------------------
# 第五步：执行与保存
# ----------------------
if __name__ == "__main__":
    # 参数设置
    input_file = "./multi_turn_dataset_2.jsonl"
    output_file = "./multi_turn_dataset_2_merge.jsonl"

    # 完整处理流程
    original_data = load_jsonl(input_file)
    cleaned_data = deduplicate(original_data)

    # 保存结果
    with open(output_file, 'w', encoding='utf-8') as f:
        for item in cleaned_data:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

    print(f"去重完成：原始数据 {len(original_data)} 条 -> 去重后 {len(cleaned_data)} 条")